# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: sungraecho
# pylint: skip-file

import os
import sys
import functools
import datetime

import fire
import numpy as np
import pandas as pd
import coin.strategy.mm.tool.archive_base as abase

from coin2.support.product_info.product_info_holder import ProductInfoHolder

sys.path.append(
  os.path.expandvars("$HOME/workspace/fastfeature/coinstrat"))

from misc import aux
LOGGER = aux.get_logger()

# Only for one product
class FeedDumperRunner(object):
  def __init__(self,
               product,
               start_ts=None,
               end_ts=None,
               product_to_sub_request=None,
               use_exchange_time=None):
    self._product = product
    self._start_ts = int(start_ts.timestamp() * 1e9)
    self._end_ts = int(end_ts.timestamp() * 1e9)
    self._bbos = []
    self._trades = []
    self.nbook = 0
    self.ntrade = 0
    self.product_to_sub_request = product_to_sub_request
    self.use_exchange_time= use_exchange_time

  def on_book_reset(self, book_builder_name, book_builder):
    product = self._product
    book_builder.subscribe(product, functools.partial(self.on_book, product))
    if hasattr(book_builder, "subscribe_trade"):
      book_builder.subscribe_trade(product,
                                   functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    timestamp = self.timestamp(book)
    if timestamp < self._start_ts or self._end_ts < timestamp:
      return
    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
      bid0_qty = book.bid0().qty
      ask0_qty = book.ask0().qty
    except (AttributeError, IndexError):
      return
    self._bbos.append({
        'timestamp': timestamp,
        'bid0_price': bid0_price,
        'ask0_price': ask0_price,
        'bid0_qty': bid0_qty,
        'ask0_qty': ask0_qty,
    })
    self.nbook += 1

  def on_trade(self, product, trade):
    timestamp = self.timestamp(trade)
    if timestamp < self._start_ts or self._end_ts < timestamp:
      return
    self._trades.append({
      'timestamp': self.timestamp(trade),
      'trade_price': trade.price,
      'trade_qty': trade.qty,
      'trade_side': trade.side,
    })
    self.ntrade += 1

  def timestamp(self, feed):
    return feed.exchange_timestamp if self.use_exchange_time else feed.timestamp

  def get_df(self, ffill=False):
    book_df = pd.DataFrame(self._bbos)
    trade_df = pd.DataFrame(self._trades)
    agg_df = book_df.merge(
      trade_df,
      on='timestamp',
      how='outer').sort_values(by='timestamp').reset_index(drop=True)
    if ffill:
      agg_df['bid0_price'] = agg_df['bid0_price'].ffill()
      agg_df['bid0_qty'] = agg_df['bid0_qty'].ffill()
      agg_df['ask0_price'] = agg_df['ask0_price'].ffill()
      agg_df['ask0_qty'] = agg_df['ask0_qty'].ffill()
    return agg_df

def get_datetimes(date, time_range):
  if '-' in time_range:
    start_time, end_time = time_range.split('-')
  else:
    start_time, end_time = '0', '24'

  def get_hhmm(hhmm_str):
    if '.' not in hhmm_str:
      hour = int(hhmm_str)
      minute = 0
    else:
      hour, minute = hhmm_str.split('.')
      hour = int(hour)
      minute = int(60 * int(minute) * 10 ** (-len(minute)))
    return datetime.timedelta(hours=hour, minutes=minute)

  date = datetime.datetime.strptime(date, "%Y%m%d")
  return date + get_hhmm(start_time), date + get_hhmm(end_time)


class FeedDumper:
  def __init__(self,
               symbol="BTC-USDT.PERPETUAL",
               date="20210907",
               time_range="14.6-16",
               exchange="Binance",
               market_type="Futures",
               recipe="v1.l1_realtime_move2bp",
               feed_machine="feed-05.ap-northeast-1.aws",
               ):
    date = str(date)
    base_flags = aux.Option({
      "trading_date": date,
      "exchange": exchange,
      "market_type": market_type,
      "time_range": time_range,
    })
    feed_flags = aux.Option({
      "feed_machine": feed_machine,
      "feed_root": "/remote/iosg/coin/data/flow",
      "worker_ids": "0",
      "use_feed_cache": True,
      "feed_cache_dir": "/remote/iosg/coin-2/feed_cache",
      "api_override": {
        f'{market_type}:{exchange}': recipe
      }
    })

    LOGGER.info("base_flags: %s", base_flags)
    LOGGER.info("feed_flags: %s", feed_flags)
    start_ts, end_ts = get_datetimes(date, time_range)
    products = abase.symbols_to_products(
      [symbol],
      base_flags
    )
    api_version_verride_map = {
      (market_type, exchange): recipe
    }
    product_to_sub_request = abase.products_to_feedsub_request(
      products, api_version_verride_map)
    runner = FeedDumperRunner(product=products[0],
                              start_ts=start_ts,
                              end_ts=end_ts,
                              product_to_sub_request=product_to_sub_request,
                              use_exchange_time=False)

    LOGGER.info("Start reading feed from archive")
    abase.run_from_archive_base(baseflags=base_flags,
                                feed_sub_flags_proto=feed_flags,
                                products=products,
                                on_book_reset=runner.on_book_reset)
    LOGGER.info("Finished reading feed from archive")
    self.runner = runner

  def dump_df(self, csv_fpath=None, ffill=False):
    if csv_fpath is None:
      csv_fpath = os.path.expandvars("$HOME/out.csv")
    df = self.runner.get_df(ffill)
    df.to_csv(csv_fpath, index=False)
    LOGGER.info("Dumped to %s", csv_fpath)

if __name__ == "__main__":
  fire.Fire(FeedDumper)

